package org.finesys.ai.service.impl;

import static org.finesys.ai.constants.EmbedConst.FILE_NAME;
import static org.finesys.ai.constants.EmbedConst.KNOWLEDGE;

import java.time.Duration;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.stream.Collectors;

import org.finesys.ai.constants.EmbedConst;
import org.finesys.ai.core.embedding.EmbeddingResponse;
import org.finesys.ai.entity.AigcDoc;
import org.finesys.ai.entity.AigcDocSlice;
import org.finesys.ai.entity.AigcModel;
import org.finesys.ai.service.AigcDocService;
import org.finesys.ai.service.AigcDocSliceService;
import org.finesys.ai.service.EmbeddingService;
import org.finesys.chat.core.base.chat.ChatRequest;
import org.finesys.chat.core.base.chat.EmbeddingR;
import org.finesys.chat.core.base.client.ChatClient;
import org.finesys.chat.core.base.document.Document;
import org.finesys.chat.core.base.document.loader.UrlDocumentLoader;
import org.finesys.chat.core.base.document.splitter.DocumentSplitter;
import org.finesys.chat.core.base.embedding.Embedding;
import org.finesys.chat.core.base.embedding.EmbeddingSearchRequest;
import org.finesys.chat.core.base.embedding.EmbeddingSearchResult;
import org.finesys.chat.core.base.embedding.EmbeddingStore;
import org.finesys.chat.core.base.embedding.factory.EmbeddingProvider;
import org.finesys.chat.core.base.embedding.filter.Filter;
import org.finesys.chat.core.base.embedding.filter.MetadataFilterBuilder;
import org.finesys.chat.core.base.segment.TextSegment;
import org.finesys.chat.core.base.service.LangEmbeddingService;
import org.finesys.common.core.util.Utils;
import org.finesys.tika.parser.ApacheTikaDocumentParser;
import org.reactivestreams.Publisher;
import org.springframework.stereotype.Service;
import org.springframework.util.StringUtils;

import com.baomidou.mybatisplus.core.toolkit.Wrappers;

import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import reactor.core.publisher.Mono;

@Service
@Slf4j
@RequiredArgsConstructor
public class EmbeddingServiceImpl implements EmbeddingService {

    private final EmbeddingProvider embeddingProvider;
    private final AigcDocService aigcDocService;
    private final AigcDocSliceService aigcDocSliceService;
    private final LangEmbeddingService langEmbeddingService;


    @Override
    @SuppressWarnings("unchecked")
    public void clearDocSlices(String docsId) {
        if (!StringUtils.hasText(docsId)) {
            return;
        }
        List<Map<String, Object>> vectorList = aigcDocSliceService.listMaps(Wrappers.<AigcDocSlice>lambdaQuery().eq(AigcDocSlice::getDocId, docsId).select(AigcDocSlice::getVectorId));
        List<String> vectorIds = vectorList.stream().filter(map -> map.get("vectorId") != null).map(map -> (String) map.get("vectorId")).collect(Collectors.toList());
        if (Utils.isNullOrEmpty(vectorIds)) {
            return;
        }
        AigcDoc doc = aigcDocService.getById(docsId);
        if (Objects.isNull(doc)) {
            return;
        }
        String knowledgeId = doc.getKnowledgeId();
        //刪除向量库中的向量
        EmbeddingStore<TextSegment> embeddingStore = embeddingProvider.getEmbeddingStore(knowledgeId);
        embeddingStore.deleteAll(vectorIds);
        //删除切片信息
        aigcDocSliceService.remove(Wrappers.<AigcDocSlice>lambdaQuery().eq(AigcDocSlice::getDocId, docsId));
        //删除文档信息
        aigcDocService.removeById(docsId);
    }

    @Override
    public void embedDocsSlice(AigcDoc data) {
        if (data == null || data.getType() == null) {
            return;
        }
        String type = data.getType();
        ChatRequest chatRequest = buildChatRequest(data);
        if (EmbedConst.INPUT.equals(type)) {
            embedInputType(data, chatRequest);
        } else {
            embedOtherType(data, chatRequest);
        }
    }

    private void embedInputType(AigcDoc data, ChatRequest chatRequest) {
        langEmbeddingService.embeddingDocs(chatRequest)
                .flatMapIterable(list -> list)
                .flatMap(embeddingR -> saveDocSlice(data, embeddingR))
                .count()
                .flatMap(sliceCount -> updateDocStatus(data.getId(), sliceCount.intValue()))
                .onErrorResume(error -> {
                    log.error("Embed input docs error", error);
                    return Mono.empty();
                });
    }


    private void embedOtherType(AigcDoc data, ChatRequest chatRequest) {
        langEmbeddingService.embeddingText(chatRequest)
                .flatMap(res -> saveDocSlice(data, res)
                        .then(updateDocStatus(data.getId(), 1)))
                .subscribe(null, error -> log.error("Embed text error", error));
    }

    private Mono<AigcDocSlice> saveDocSlice(AigcDoc data, EmbeddingR result) {
        AigcDocSlice docSlice = new AigcDocSlice()
                .setDocId(data.getId())
                .setVectorId(result.getVectorId())
                .setKnowledgeId(data.getKnowledgeId())
                .setName(data.getName())
                .setContent(result.getText())
                .setWordNum(result.getText().length())
                .setStatus(true);
        aigcDocSliceService.save(docSlice);
        return Mono.just(docSlice);
    }

    private Mono<Void> updateDocStatus(String docId, int sliceNumber) {
        aigcDocService.updateById(
                new AigcDoc()
                        .setId(docId)
                        .setSliceNumber(sliceNumber)
                        .setSliceStatus(true));
        return Mono.empty();
    }

    @Override
    public List<Map<String, Object>> search(AigcDoc data) {
        if (data == null || !StringUtils.hasText(data.getKnowledgeId()) || !StringUtils.hasText(data.getContent())) {
            return Collections.emptyList();
        }
        if (!StringUtils.hasText(data.getType())) {
            data.setType(EmbedConst.INPUT);
        }
        String knowledgeId = data.getKnowledgeId();
        //刪除向量库中的向量
        EmbeddingStore<TextSegment> embeddingStore = embeddingProvider.getEmbeddingStore(knowledgeId);
        ChatClient chatClient = embeddingProvider.getChatClient(knowledgeId);
        Filter filter = MetadataFilterBuilder.of(KNOWLEDGE).equal(knowledgeId);

        ChatRequest chatRequest = buildChatRequest(data);
        Publisher<EmbeddingResponse> embeddingResponsePublisher = chatClient.embedding(chatRequest);
        Embedding embedding = Mono.from(embeddingResponsePublisher).map(res -> Embedding.from(res.embedding())).block(Duration.ofSeconds(30));
        EmbeddingSearchRequest embeddingSearchRequest = EmbeddingSearchRequest.builder()
                .queryEmbedding(embedding)
                .filter(filter).build();
        EmbeddingSearchResult<TextSegment> list = embeddingStore.search(embeddingSearchRequest);
        List<Map<String, Object>> result = new ArrayList<>();
        list.getMatches().forEach(i -> {
            TextSegment embedded = i.getEmbeded();
            Map<String, Object> map = embedded.getMetadata().toMap();
            map.put("text", embedded.getText());
            result.add(map);
        });
        return result;
    }


    private ChatRequest buildChatRequest(AigcDoc data) {
        //获取模型信息
        AigcModel aigcModel = embeddingProvider.getAigcModel(data.getKnowledgeId());
        //构建输入列表
        List<String> inputs = new ArrayList<>();
        if (EmbedConst.INPUT.equals(data.getType())) {
            inputs.add(data.getContent());
        } else {
            //解析上传的文档
            Document document = UrlDocumentLoader.load(data.getUrl(), new ApacheTikaDocumentParser());
            document.getMetadata().put(KNOWLEDGE, data.getKnowledgeId())
                    .put(FILE_NAME, data.getName());
            //根据URL获取内容并拆分为段落
            DocumentSplitter splitter = EmbeddingProvider.splitter();
            List<TextSegment> textSegments = splitter.split(document);
            //转换为输入列表
            textSegments.stream().map(TextSegment::getText).forEach(inputs::add);
        }
        //构建请求
        return new ChatRequest()
                .setDocId(data.getId())
                .setKnowledgeId(data.getKnowledgeId())
                .setModel(aigcModel.getModelName())
                .setDimensions(aigcModel.getDimension())
                .setInput(inputs)
                .setTemperature(aigcModel.getTemperature())
                .setTopP(aigcModel.getTopP())
                .setImageQuality(aigcModel.getImageQuality())
                .setImageSize(aigcModel.getImageSize())
                .setImageStyle(aigcModel.getImageStyle())
                ;
    }
}
